# https://pypi.org/project/KunQuant/
import numpy as np
import pandas as pd
from KunQuant.jit import cfake
from KunQuant.Driver import KunCompilerConfig
from KunQuant.Op import Builder, Input, Output
from KunQuant.Stage import Function
from KunQuant.predefined import Alpha101
from KunQuant.runner import KunRunner as kr
from KunQuant.predefined.Alpha101 import AllData, all_alpha

NUM_STOCKS=88

def test1():
  builder = Builder()
  with builder:
      vclose = Input("close")
      low = Input("low")
      high = Input("high")
      vopen = Input("open")
      amount = Input("amount")
      vol = Input("volume")
      all_data = Alpha101.AllData(low=low,high=high,close=vclose,open=vopen, amount=amount, volume=vol)
      Output(Alpha101.alpha001(all_data), "alpha001")
  f = Function(builder.ops)

  print(f)


np.random.seed(42)
num_time=261
def generate_1_simulated_market_data(
    start_date='2025-01-01',
    end_date='2025-12-31',
    base_price=58.0,
    volatility=0.02,  # 日波动率（标准差）
    volume_mean=1000000,  # 平均成交量
    volume_std=300000     # 成交量标准差
):
    """
    生成模拟的股票行情数据。

    参数:
    - symbol: 股票代码
    - start_date: 开始日期 (字符串格式 'YYYY-MM-DD')
    - end_date: 结束日期 (字符串格式 'YYYY-MM-DD')
    - base_price: 基础价格（起始价格）
    - volatility: 价格波动率（日收益率的标准差）
    - volume_mean: 平均成交量
    - volume_std: 成交量标准差

    返回:
    - 包含模拟行情数据的 pandas DataFrame
    """
    
    # 1. 生成交易日（排除周末）
    dates = pd.date_range(start=start_date, end=end_date, freq='B')  # 'B' 表示工作日
    num_days = len(dates)
    
    # 2. 生成收盘价 (使用几何布朗运动近似)
    # 日收益率 = 随机正态分布 * 波动率
    daily_returns = np.random.normal(loc=0, scale=volatility, size=num_days)
    # 累积收益率
    price_path = base_price * np.exp(np.cumsum(daily_returns))
    
    # 3. 生成开盘价 (假设开盘价接近前一日收盘价，但有小波动)
    # 第一天的开盘价设为基础价格
    open_price = np.zeros(num_days)
    open_price[0] = base_price
    # 后续开盘价为前一日收盘价加上小的随机波动
    for i in range(1, num_days):
        open_price[i] = price_path[i-1] * (1 + np.random.normal(0, volatility * 0.3))
    
    # 4. 生成最高价和最低价
    # 假设当日价格在开盘价和收盘价之间波动，并有一定范围
    high_price = np.maximum(open_price, price_path)  # 至少是开盘或收盘的最高值
    low_price = np.minimum(open_price, price_path)   # 至少是开盘或收盘的最低值
    
    # 添加随机波动来扩大高低价范围
    # 波动幅度与当日价格水平和波动率相关
    price_levels = (open_price + price_path) / 2
    fluctuation = price_levels * volatility * np.random.uniform(0.5, 2, size=num_days)
    
    high_price += fluctuation
    low_price -= fluctuation
    
    # 确保最低价 > 0
    low_price = np.maximum(low_price, price_levels * 0.8)  # 限制最低价不低于某个比例
    
    # 5. 生成成交量 (正态分布，但取绝对值确保为正)
    volume = np.abs(np.random.normal(loc=volume_mean, scale=volume_std, size=num_days)).astype(int)
    
    # 6. 构建 DataFrame
    df = pd.DataFrame({
        'date': dates,
        'open': open_price,
        'high': high_price,
        'low': low_price,
        'close': price_path,
        'volume': volume
    })
    
    # 设置日期为索引（可选）
    df.set_index('date', inplace=True)
    
    # 四舍五入到合适的小数位
    price_columns = ['open', 'high', 'low', 'close']
    df[price_columns] = df[price_columns].round(2)
    df['amount'] = df['close']*df['volume']
    num_time=len(df)
    return df

def generate_simulated_market_datas(num_stocks=NUM_STOCKS):
  dfs = []
  for i in range(0, num_stocks):      
    dfs.append(generate_1_simulated_market_data( base_price = np.random.normal(58, 18)))
  return dfs


def get_data():
  dfs = generate_simulated_market_datas()
  # print(dfs)
  cols = dfs[0].columns.values
  col2idx = dict(zip(cols, range(len(cols))))
  print("columns to index", col2idx)
  num_time = len(dfs[0])
  print("dimension in time", num_time)

  collected = np.empty((len(col2idx), NUM_STOCKS, len(dfs[0])), dtype="float32")
  for stockidx, data in enumerate(dfs):
      for colname, colidx in col2idx.items():
          mat = data[colname].to_numpy()
          collected[colidx, stockidx, :] = mat
  return(collected)
 

def test2():
  collected=get_data()  
  print(collected)
  
def get_transposed_data():
  collected=get_data()  
  # Transpose the matrix to [features, time, stocks]

  # [features, stocks, time] => [features, time, stocks]
  transposed = collected.transpose((0, 2, 1))
  transposed = np.ascontiguousarray(transposed)
  print(transposed)

  # Now fill the input data in a dict of {"open": matrix_open, "close": ...}

  input_dict = dict()
  col2idx={'open': 0, 'high': 1, 'low': 2, 'close': 3, 'volume': 4, 'amount': 5}
  for colname, colidx in col2idx.items():
      input_dict[colname] = transposed[colidx]
  return input_dict


def test3():
  data = get_transposed_data()
  print(data)

def test4():
  input_dict = get_transposed_data()

  # using 4 threads
  executor = kr.createMultiThreadExecutor(4)
  builder = Builder()
  with builder:
      vclose = Input("close")
      low = Input("low")
      high = Input("high")
      vopen = Input("open")
      amount = Input("amount")
      vol = Input("volume")
      all_data = Alpha101.AllData(low=low,high=high,close=vclose,open=vopen, amount=amount, volume=vol)
      # Output(Alpha101.alpha001(all_data), "alpha001")
      for f in all_alpha:
          out = f(all_data)
          Output(out, f.__name__)

  f = Function(builder.ops)
  lib = cfake.compileit([("alpha101", f, KunCompilerConfig(input_layout="TS", output_layout="TS"))], "out_first_lib", cfake.CppCompilerConfig())
  modu = lib.getModule("alpha101")
  out = kr.runGraph(executor, modu, input_dict, 0, num_time)
  # print("Shape of alpha101", out["alpha001"].shape)
  # print("Result of alpha101", out["alpha001"])
  #   
  # df = pd.DataFrame(out["alpha001"])
  # print(df)

  # for f in all_alpha:
  #   sq = f.__name__
  #   df = pd.DataFrame(out[sq])
  #   print(sq)    
  #   print(df)

  for sq in out:
    df = pd.DataFrame(out[sq])
    print(sq)    
    print(df)


if __name__ == "__main__":
  #  test1()

  #  df1 = generate_1_simulated_market_data()
  #  print(df1)
  # test2()
  # test3()
  test4()